import os
import numpy as np
from tqdm import tqdm
from pprint import pprint

from torch.utils.data import Dataset, DataLoader
from wputils.utils.resize import m3d, m3dc, m3dc_inv, r3d, r3d0
from wputils.utils.norm import nml
from wputils.utils.image import rbbox, bshape, bcenter
from wpdata.dloader.utils import closest_divisible, one_hot

from wputils.utils.io import rnii, wnii, rnpz, wnpy, rnpy, cdirs, fdirs, gsiz, gmdf, wseq, rseq


def acquire_img(root):
    fs = [os.path.join(root, f) for f in os.listdir(root) if f.endswith('.nii.gz')]
    if len(fs) == 1:
        return fs[0]
    else:
        sizs = [gsiz(f) for f in fs]
        return fs[sizs.index(max(sizs))], fs[sizs.index(min(sizs))]


if __name__ == '__main__':
    #  sorting files:

    # -------------------------------- writting ---------------------------

    # ROOT_PATH = ['/share_data/liupan/HCPAPARC-SORT/npyfile/crop',
    #              '/share_data/liupan/HCPAPARC-SORT/npyfile/rsiz']
    #
    # files = fdirs(ROOT_PATH, 'npy')

    files = rseq('/share_data/hematoma_nnunet/24h_manifest.noMoreThan32Slices.json')
    print(len(files.get('trainfiles')), len(files.get('testfiles')))
    target_dir = '/share_data/liupan/hematoma/npyfiles/24h/valid'
    for f in tqdm(files.get('testfiles')):
        imgpath = f.get('image')
        labpath = f.get('label')
        imgarr, stk = rnii(imgpath)
        labarr, ___ = rnii(labpath)
        if imgarr.shape[0] < 16:
            continue
        imgarr = m3d(imgarr.astype(np.int32), (24, 384, 384))
        labarr = m3d(labarr.astype(np.int8), (24, 384, 384))
        imgarr = np.clip(imgarr, -1024, 2048) + 1024
        imgarr = imgarr.astype(np.uint16)
        labarr = labarr.astype(bool).astype(np.uint8)
        fmdf = gmdf(imgpath)
        dst_npz_id = os.path.join(target_dir, fmdf)
        id_img_path = os.path.join(dst_npz_id, 'img.npy')
        id_lab_path = os.path.join(dst_npz_id, 'lab.npy')
        wnpy(imgarr, id_img_path)
        wnpy(labarr, id_lab_path)

    # for file in files:
    #     imgpath = os.path.join(file, 'img.npy')
    #     labpath = os.path.join(file, 'lab.npy')
    #
    #     imgarr = rnpy(imgpath)
    #     labarr = rnpy(labpath)
    #
    #     print(np.amin(imgarr), np.amax(imgarr), np.amin(labarr), np.amax(labarr))

    # --------------------------------- writting ---------------------

    # ROOT_PATHS = ['/share_data/liupan/HCPAPARC-SORT/nifti/test']
    # ROOT_PATH = '/share_data/liupan/HCPAPARC-SORT/nifti'
    # files = fdirs(ROOT_PATHS, 'nii.gz')
    #
    # print(len(files))
    # print(files[0])
    #
    # for filepath in tqdm(files):
    #     imgpth, labpath = acquire_img(filepath)
    #     fmdf = gmdf(imgpth)
    #     imgarr, stk = rnii(imgpth)
    #     labarr, _ = rnii(labpath)
    #
    #     imgarr = np.clip(imgarr, 0, 65535)
    #
    #     dim = imgarr.shape
    #
    #     dst_npz_id = ROOT_PATH.replace('nifti', 'npyfile/identity/valid')
    #     id_img_path = os.path.join(dst_npz_id, fmdf, 'img.npy')
    #     id_lab_path = os.path.join(dst_npz_id, fmdf, 'lab.npy')
    #     id_bny_path = os.path.join(dst_npz_id, fmdf, 'bny.npy')
    #
    #     wnpy(imgarr.astype(np.uint16), id_img_path)
    #     wnpy(labarr.astype(np.uint8), id_lab_path)
    #     wnpy(labarr.astype(bool).astype(np.uint8), id_bny_path)
    #
    #     bbox = rbbox(labarr)
    #     bbshape, bbc = bshape(bbox), bcenter(bbox)
    #
    #     rshape = tuple([closest_divisible(i, 8) for i in bbshape])
    #
    #     rimg, rlab = m3dc(imgarr, bbc, dim=rshape, bg=0), m3dc(labarr, bbc, dim=rshape, bg=0)
    #
    #     tdim = (152, 192, 136)
    #
    #     mimg, mlab = m3d(rimg.copy(), tdim, 0).astype(np.uint16), m3d(rlab.copy(), tdim, 0).astype(np.uint8)
    #     cimg, clab = r3d(rimg.copy(), tdim, 1).astype(np.uint16), r3d0(rlab.copy(), tdim).astype(np.uint8)
    #
    #     wnpy(mimg.astype(np.uint16), id_img_path.replace('identity', 'crop'))
    #     wnpy(mlab.astype(np.uint8), id_lab_path.replace('identity', 'crop'))
    #     wnpy(mlab.astype(bool).astype(np.uint8), id_bny_path.replace('identity', 'crop'))
    #
    #     wnpy(cimg.astype(np.uint16), id_img_path.replace('identity', 'rsiz'))
    #     wnpy(clab.astype(np.uint8), id_lab_path.replace('identity', 'rsiz'))
    #     wnpy(clab.astype(bool).astype(np.uint8), id_bny_path.replace('identity', 'rsiz'))


    # -------------------------------- checking ---------------------------

    # ROOT_PATH = ['/share_data/liupan/HCPAPARC-SORT/npyfile/crop',
    #              '/share_data/liupan/HCPAPARC-SORT/npyfile/rsiz']
    #
    # files = fdirs(ROOT_PATH, 'npy')
    #
    # for file in files:
    #     imgpath = os.path.join(file, 'img.npy')
    #     labpath = os.path.join(file, 'lab.npy')
    #
    #     imgarr = rnpy(imgpath)
    #     labarr = rnpy(labpath)
    #
    #     print(np.amin(imgarr), np.amax(imgarr), np.amin(labarr), np.amax(labarr))

    # --------------------------------- sortng the scores -----------------------

    # JSON_PATH = '/data/liupan/WingsofPanda/WP-Head/wphead/t1HCPAPARC/nets/unet/checkpoint/0_warmup_crop_rsiz/e126_vl_0.397986_avl_-1.json'
    #
    #
    # def sort_dict_by_values(input_dict: dict) -> list:
    #     sorted_list = [f"{k}_score_{v}" for k, v in sorted(input_dict.items(), key=lambda item: item[1])]
    #     return sorted_list
    #
    # scores = rseq(JSON_PATH)
    #
    # sortedd = sort_dict_by_values(scores.get('sf'))
    # pprint(sortedd)
